/*******************************************************************************
 * *
 * *             COPYRIGHT 2012 CLEARLEAP, INC. ALL RIGHTS RESERVED.
 * *             CLEARLEAP CONFIDENTIAL PROPRIETARY
 * *             TEMPLATE VERSION R01.00
 * *
 * ********************************************************************************
 * *
 * *  FILE NAME: XmlDigitalSignatureSigner.java
 * *
 * *---------------------------------- PURPOSE ------------------------------------
 * *
 * * Utilities to sign XML document, used to sign SAML Requests before sending
 * *
 * *
 * *----------------------------- REVISION HISTORY --------------------------------
 * *
 * * Date      ID        Tracking#      Description
 * * --------  --------  -------------  -------------------------------------------
 * * 01.07                              Created
 * *
 * *--------------------------- INCLUDES ----------------------------*
 * *
 ******************************************************************************/
/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package com.easysale.util;

import java.io.ByteArrayInputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.StringWriter;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.cert.CertificateFactory;
import java.util.Collections;
import java.util.logging.Logger;

import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathFactory;

import org.w3c.dom.Document;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

/**
 *
 * @author minhtran
 */
public class XmlDigitalSignatureSigner {

    protected static Logger logger = Logger
            .getLogger("XmlDigitalSignatureSigner");

    public static String signAuthnXML(String xmlString, String certFile,
            String keyStoreFile) throws Exception {
        Document XML = getDocument(xmlString);

        return signXML(XML, certFile, keyStoreFile, XML.getDocumentElement(),
                XML.getDocumentElement().getFirstChild().getNextSibling()
                        .getNextSibling());
    }

    public static String signAuthzXML(String xmlString, String certFile,
            String keyStoreFile) throws Exception {
        Document XML = getDocument(xmlString);

        XPath xpath = XPathFactory.newInstance().newXPath();
        xpath.setNamespaceContext(new HardcodedNamespaceResolver());
        NodeList parentNode = (NodeList) xpath
                .evaluate(
                        "/soap11:Envelope/soap11:Body/xacml-samlp:XACMLAuthzDecisionQuery",
                        XML.getDocumentElement(), XPathConstants.NODESET);
        NodeList nextSiblingNode = (NodeList) xpath
                .evaluate(
                        "/soap11:Envelope/soap11:Body/xacml-samlp:XACMLAuthzDecisionQuery/xacml-context:Request",
                        XML.getDocumentElement(), XPathConstants.NODESET);

        return signXML(XML, certFile, keyStoreFile, parentNode.item(0),
                nextSiblingNode.item(0));
    }

    private static Document getDocument(String xmlString) throws SAXException,
            IOException, ParserConfigurationException {
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        return dbf.newDocumentBuilder().parse(
                new ByteArrayInputStream(xmlString.getBytes()));
    }

    /**
     * Sign xml private for 2 methods sign request and sign authorization.
     *
     * @param xml
     * @param certFile
     * @param keyStoreFile
     * @param parent
     * @param nextSibling
     * @return
     * @throws Exception
     */
    private static String signXML(Document xml, String certFile,
            String keyStoreFile, Node parent, Node nextSibling)
            throws Exception {
        XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");
        Reference ref = fac.newReference("",
                fac.newDigestMethod(DigestMethod.SHA1, null));
        SignedInfo si = fac
                .newSignedInfo(fac.newCanonicalizationMethod(
                        CanonicalizationMethod.INCLUSIVE_WITH_COMMENTS,
                        (C14NMethodParameterSpec) null), fac
                        .newSignatureMethod(SignatureMethod.DSA_SHA1, null),
                        Collections.singletonList(ref));

        // Prepare keystore
        PrivateKey privateKey = null;
        KeyPair kp = KeyStoreInfo.getKeyPair(keyStoreFile, "sp1234", "kp1234",
                "biz");

        privateKey = kp.getPrivate();
        KeyInfo ki = null;

        KeyInfoFactory kif = fac.getKeyInfoFactory();

        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        FileInputStream fis = new FileInputStream(certFile);
        java.security.cert.Certificate cert = cf.generateCertificate(fis);
        fis.close();
        X509Data x509d = kif.newX509Data(Collections.singletonList(cert));
        ki = kif.newKeyInfo(Collections.singletonList(x509d));

        XMLSignature signature = fac.newXMLSignature(si, ki);
        DOMSignContext dsc = new DOMSignContext(privateKey, parent, nextSibling);
        dsc.putNamespacePrefix(XMLSignature.XMLNS, "ds");

        signature.sign(dsc);

        // Materialize into an xml document
        StringWriter stringWriter = new StringWriter();
        TransformerFactory tf = TransformerFactory.newInstance();
        Transformer trans = tf.newTransformer();
        trans.transform(new DOMSource(xml), new StreamResult(stringWriter));
        String strFileContent = stringWriter.toString();

        return strFileContent;
    }
}
